package org.example.api.tableapi.udf;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.java.StreamTableEnvironment;
import org.apache.flink.table.functions.AggregateFunction;
import org.apache.flink.types.Row;
import org.example.api.bean.SensorReading;

/**
 * @author huangqihan
 * @date 2021/3/10
 */
public class AggFunc {

    public static void main(String[] args) throws Exception {
        // 创建执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        DataStreamSource<String> inputStream = env.readTextFile("src/main/resources/sensor.txt");

        SingleOutputStreamOperator<SensorReading> dataStream = inputStream.map(line -> {
            String[] fields = line.split(",");
            return new SensorReading(fields[0], new Long(fields[1]), new Double(fields[2]));
        });

        // 创建表环境
        StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);

        // 基于数据流创建表
        Table table = tableEnv.fromDataStream(dataStream, "id, timestamp as ts, temperature as temp ");

        // 4
        // 4.1 table api
        Acc acc = new Acc();
        tableEnv.registerFunction("acc", acc );
        Table res = table.groupBy("id").aggregate("acc(temp) as avgTemp")
                .select("id, avgTemp");

        // 4.2 sql
        tableEnv.createTemporaryView("sensor", table);
        Table sqlRes = tableEnv.sqlQuery("select id, acc(temp) from sensor group by id");

        tableEnv.toRetractStream(res, Row.class).print("res");
        tableEnv.toRetractStream(sqlRes, Row.class).print("sqlRes");

        env.execute();

    }

    public static class Acc extends AggregateFunction<Double, Tuple2<Double, Integer>> {

        @Override
        public Double getValue(Tuple2<Double, Integer> tuple) {
            return tuple.f0 / tuple.f1;
        }

        @Override
        public Tuple2<Double, Integer> createAccumulator() {
            return new Tuple2<>(0.0, 0);
        }

        // 必须实现 accumulate 方法
        public void accumulate(Tuple2<Double, Integer> acc, Double temp) {
            acc.f0 += temp;
            acc.f1++;
        }
    }
}
